package srt.tool;

import java.util.ArrayList;
import java.util.List;

import srt.ast.AssertStmt;
import srt.ast.AssignStmt;
import srt.ast.AssumeStmt;
import srt.ast.BinaryExpr;
import srt.ast.BlockStmt;
import srt.ast.Decl;
import srt.ast.DeclRef;
import srt.ast.Expr;
import srt.ast.HavocStmt;
import srt.ast.IfStmt;
import srt.ast.IntLiteral;
import srt.ast.Program;
import srt.ast.Stmt;
import srt.ast.TernaryExpr;
import srt.ast.UnaryExpr;
import srt.ast.visitor.impl.DefaultVisitor;
//Visitor implementation for predication 
public class PredicationVisitor extends DefaultVisitor {

	int freshPredicateNum = 0; //fresh Predicate var name 
	int freshVariableNum = 0; //fres var name num

	//Global predicate
	DeclRef globalPredicate = new DeclRef("G");

	DeclRef parentPred = null;

	public PredicationVisitor() {
		super(true);
	}
	//return $A0, $A1, ../ 
	private String getFreshPredicateName() {
		freshPredicateNum++;
		return "$A" + freshPredicateNum;
	}

	//return fresh0, fresh1, ../
	private String getFreshVariableName() {
		freshVariableNum++;
		return "fresh" + freshVariableNum;
	}

	//Predication for if statements
	@Override
	public Object visit(IfStmt ifStmt) {

		Expr thenExpr;
		Expr elseExpr;

		List<Stmt> statements = new ArrayList<Stmt>();

		DeclRef predicateP = new DeclRef(getFreshPredicateName());

		if (parentPred != null) {
			thenExpr = new BinaryExpr(BinaryExpr.LAND, parentPred,
					ifStmt.getCondition());
		} else {
			thenExpr = ifStmt.getCondition();
		}

		statements.add(new AssignStmt(predicateP, thenExpr));

		DeclRef predicateQ = null;

		if (ifStmt.getElseStmt() != null) {

			predicateQ = new DeclRef(getFreshPredicateName());

			if (parentPred != null) {
				elseExpr = new BinaryExpr(BinaryExpr.LAND, parentPred,
						new UnaryExpr(UnaryExpr.LNOT, ifStmt.getCondition()));
			} else {
				elseExpr = new UnaryExpr(UnaryExpr.LNOT, ifStmt.getCondition());
			}

			statements.add(new AssignStmt(predicateQ, elseExpr));
		}

		DeclRef oldParentPred = parentPred;
		this.parentPred = predicateP;

		Stmt elseStmt, thenStmt;
		thenStmt = (Stmt) visit(ifStmt.getThenStmt());
		statements.add(thenStmt);

		if (ifStmt.getElseStmt() != null) {
			this.parentPred = predicateQ;
			elseStmt = (Stmt) visit(ifStmt.getElseStmt());
			statements.add(elseStmt);
		}

		parentPred = oldParentPred;

		return new BlockStmt(statements, ifStmt.getNodeInfo());
	}

	//Predication for assertments
	@Override
	public Object visit(AssertStmt assertStmt) {

		if (this.parentPred != null) {
			// (G && P) => E
			return new AssertStmt(new BinaryExpr(BinaryExpr.LOR, new UnaryExpr(
					UnaryExpr.LNOT, new BinaryExpr(BinaryExpr.LAND,
							globalPredicate, parentPred)),
					assertStmt.getCondition()), assertStmt.getNodeInfo());
		} else {
			// (G && true) => E
			return new AssertStmt(
					new BinaryExpr(BinaryExpr.LOR, new UnaryExpr(
							UnaryExpr.LNOT, globalPredicate),
							assertStmt.getCondition()),
					assertStmt.getNodeInfo());
		}

	}

	//Predication for assignmetns
	@Override
	public Object visit(AssignStmt assignment) {

		if (this.parentPred != null) {
			TernaryExpr exp = new TernaryExpr(new BinaryExpr(BinaryExpr.LAND,
					globalPredicate, parentPred), assignment.getRhs(),
					assignment.getLhs());
			return new AssignStmt(assignment.getLhs(), exp,
					assignment.getNodeInfo());
		}

		return super.visit(assignment);
	}

	@Override
	public Object visit(AssumeStmt assumeStmt) {

		if (parentPred != null) {
			// G = G && (P => E)
			return new AssignStmt(globalPredicate, new BinaryExpr(
					BinaryExpr.LAND, globalPredicate, new BinaryExpr(
							BinaryExpr.LOR, new UnaryExpr(UnaryExpr.LNOT,
									parentPred), assumeStmt.getCondition())));

		} else {
			// G = G && E
			return new AssignStmt(globalPredicate,
					new BinaryExpr(BinaryExpr.LAND, globalPredicate,
							assumeStmt.getCondition()));
		}
	}
	
	//predication for havoc
	@Override
	public Object visit(HavocStmt havocStmt) {

		// x = ((G && P) ? h : x) where h is fresh

		if (parentPred != null) {
			return new AssignStmt(havocStmt.getVariable(),
					new TernaryExpr(new BinaryExpr(BinaryExpr.LAND,
							globalPredicate, parentPred), new DeclRef(
							getFreshVariableName()), havocStmt.getVariable()),
					havocStmt.getNodeInfo());
		} else {
			return new AssignStmt(havocStmt.getVariable(), new TernaryExpr(
					globalPredicate, new DeclRef(getFreshVariableName()),
					havocStmt.getVariable()), havocStmt.getNodeInfo());
		}
	}

	//initiation of G
	@Override
	public Object visit(Program program) {
		Decl globalPredDeclaration = new Decl(globalPredicate.getName(), "int");
		Stmt globalPredAssignment = new AssignStmt(globalPredicate,
				new IntLiteral(1));
		Stmt oldBlock = (Stmt) visit(program.getBlockStmt());
		BlockStmt newBlock = new BlockStmt(new Stmt[] { globalPredDeclaration,
				globalPredAssignment, oldBlock });
		return new Program(program.getFunctionName(), program.getDeclList(),
				newBlock, program.getNodeInfo());
	}

}